{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utility Classes to support JSON generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "class XaccJSONEncoder(json.JSONEncoder):\n",
    "    \"\"\"\n",
    "    JSON encoder for NumPy arrays and complex numbers.\n",
    "    This functions as the standard JSON Encoder but adds support\n",
    "    for encoding:\n",
    "        complex numbers z as lists [z.real, z.imag]\n",
    "        numpy.ndarrays as nested lists.\n",
    "    \"\"\"\n",
    "    def default(self, obj):\n",
    "        if isinstance(obj, np.ndarray):\n",
    "            return obj.tolist()\n",
    "        if isinstance(obj, complex):\n",
    "            return [obj.real, obj.imag]\n",
    "        if hasattr(obj, \"to_dict\"):\n",
    "            return obj.to_dict()\n",
    "        return super().default(obj)\n",
    "\n",
    "def clean_nones(value):\n",
    "    \"\"\"\n",
    "    Recursively remove all None values from dictionaries and lists, and returns\n",
    "    the result as a new dictionary or list.\n",
    "    \"\"\"\n",
    "    if isinstance(value, list):\n",
    "        return [clean_nones(x) for x in value if x is not None]\n",
    "    elif isinstance(value, dict):\n",
    "        return {\n",
    "            key: clean_nones(val)\n",
    "            for key, val in value.items()\n",
    "            if val is not None\n",
    "        }\n",
    "    else:\n",
    "        return value    \n",
    "    \n",
    "def json_serialize(obj):\n",
    "    return json.loads(json.dumps(clean_nones(obj), cls=XaccJSONEncoder))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ChannelKrausOp:\n",
    "    '''\n",
    "    Represents a noise channel in terms of Kraus operators: \n",
    "    Args:\n",
    "        noise_qubits (List[String]): qubits that this operator acts on (If none, applying on gate qubits)\n",
    "        matrix: the list of Kraus ops (as square matrices) representing a noise channel.\n",
    "        Note: must satisfy the CPTP condition\n",
    "    '''\n",
    "    def __init__(self, matrix, noise_qubits = None):\n",
    "        self.matrix = matrix\n",
    "        self.noise_qubits = noise_qubits\n",
    "    \n",
    "    \"\"\"\n",
    "    Returns:\n",
    "        dict: a (JSON) dictionary for a KrausOp\n",
    "    \"\"\"\n",
    "    def to_dict(self):\n",
    "        kraus = {\n",
    "            \"matrix\": list(self.matrix),\n",
    "            \"noise_qubits\": self.noise_qubits\n",
    "        }\n",
    "        return json_serialize(kraus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GateNoise:\n",
    "    '''\n",
    "    Represents noise ops (as a list of Kraus operators) associated with a quantum gate\n",
    "    Args:\n",
    "        name (string): gate name\n",
    "        register_location (List[String]): gate qubits (as a list of register labels)\n",
    "        noise_channels (List[KrausOp]): list of noise channels\n",
    "    '''    \n",
    "    def __init__(self, name, register_location, noise_kraus_ops):\n",
    "        self.name = name\n",
    "        self.register_location = register_location\n",
    "        self.noise_kraus_ops = noise_kraus_ops\n",
    "    \n",
    "    def to_dict(self):\n",
    "        gate_noise = {\n",
    "            \"gate_name\": self.name,\n",
    "            \"register_location\": self.register_location,\n",
    "            \"noise_channels\": self.noise_kraus_ops\n",
    "        }\n",
    "        return json_serialize(gate_noise)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReadoutError:\n",
    "    '''\n",
    "    Represents the error in measuring qubits.\n",
    "    We use two error probabilities: prob_meas0_prep1 (P(0|1)) and prob_meas1_prep0 (P(1|0))\n",
    "    to capture the ReadoutError.\n",
    "    Note: these values will be fed to the simulators, which may use different simulation methods. \n",
    "    Args:\n",
    "        register_location (String): qubit label\n",
    "        prob_meas0_prep1: Prob(0|1)\n",
    "        prob_meas1_prep0: Prob(1|0)\n",
    "    '''    \n",
    "    def __init__(self, register_location, prob_meas0_prep1, prob_meas1_prep0):\n",
    "        self.register_location = register_location\n",
    "        self.prob_meas0_prep1 = prob_meas0_prep1\n",
    "        self.prob_meas1_prep0 = prob_meas1_prep0\n",
    "\n",
    "    def to_dict(self):\n",
    "        gate_noise = {\n",
    "            \"register_location\": self.register_location,\n",
    "            \"prob_meas0_prep1\": self.prob_meas0_prep1,\n",
    "            \"prob_meas1_prep0\": self.prob_meas1_prep0\n",
    "        }\n",
    "        return json_serialize(gate_noise)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NoiseModel:\n",
    "    '''\n",
    "    Represents a noise model as a list of gate noise specifications\n",
    "    Args:\n",
    "        bit_order: MSB or LSB \n",
    "        gate_noise (List[GateNoise]): list of gate noise specifications (for Kraus ops look-up)\n",
    "        readout_errors (List[ReadoutError]): list of readout errors\n",
    "    '''    \n",
    "    def __init__(self, gate_noise, bit_order = 'MSB', readout_errors = None):\n",
    "        self.gate_noise = gate_noise\n",
    "        self.bit_order = bit_order\n",
    "        self.readout_errors = readout_errors\n",
    "    def to_dict(self):\n",
    "        noise_model = {\n",
    "            \"gate_noise\": self.gate_noise,\n",
    "            \"bit_order\": self.bit_order,\n",
    "            \"readout_errors\": self.readout_errors\n",
    "        }\n",
    "        return json_serialize(noise_model)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 1: single-qubit depolarization\n",
    "\n",
    "In this example, we show how to devise a depolarization noise channel and generate a noise model JSON for that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "# primative gates and basic operators\n",
    "sx = np.array([[0,1],[1,0]])\n",
    "sy = np.array([[0,-1j],[1j,0]])\n",
    "sz = np.array([[1,0],[0,-1]])\n",
    "s0 = np.eye(2)\n",
    "# expand depolarizing channel in the Krauss basis\n",
    "EDP = lambda l: np.array((np.sqrt(1-l)*s0, np.sqrt(l/3.0)*sx, np.sqrt(l/3.0)*sy, np.sqrt(l/3.0)*sz)) \n",
    "# Trace preserving condition. \n",
    "# Check it -> should return identity matrix\n",
    "TP = lambda l: np.tensordot(EDP(l).swapaxes(1,2),EDP(l),axes=((0,2),(0,2))) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "# depolarizing rate\n",
    "gamma = 0.01                                        \n",
    "depol = EDP(gamma)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Data for validation\n",
    "\n",
    "We compute analytically the density matrix for a 'noisy' X gate having this depolarizing channel.\n",
    "We generate a simple noise model that only contains noise info for X gate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a gate noise instance for X gate on qubit \"0\"\n",
    "# having that depolarizing channel.\n",
    "xNoise = GateNoise(\"X\", [\"0\"], [ChannelKrausOp(depol)])\n",
    "# The noise model just contains a single entry for now.\n",
    "noiseModel = NoiseModel([xNoise])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'{\"gate_noise\": [{\"gate_name\": \"X\", \"register_location\": [\"0\"], \"noise_channels\": [{\"matrix\": [[[[0.99498743710662, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.99498743710662, 0.0]]], [[[0.0, 0.0], [0.05773502691896258, 0.0]], [[0.05773502691896258, 0.0], [0.0, 0.0]]], [[[0.0, 0.0], [0.0, -0.05773502691896258]], [[0.0, 0.05773502691896258], [0.0, 0.0]]], [[[0.05773502691896258, 0.0], [0.0, 0.0]], [[0.0, 0.0], [-0.05773502691896258, 0.0]]]]}]}], \"bit_order\": \"MSB\"}'"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# This JSON can be used with XACC to initialize noisy simulation.\n",
    "json.dumps(noiseModel.to_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Validation result:\n",
    "Applying the depolarizing channel on |1><1| density matrix (result of a perfect X gate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0.00666667+0.j 0.        +0.j]\n",
      " [0.        +0.j 0.99333333+0.j]]\n"
     ]
    }
   ],
   "source": [
    "# rho = |1><1|\n",
    "rho = np.array([[0,0],[0,1]])\n",
    "result = np.zeros((2,2), dtype=complex)\n",
    "for kraus in depol:\n",
    "   result +=  kraus@rho@kraus.conj().T\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 2: Amplitude damping channel\n",
    "\n",
    "A single-qubit amplitude damping channel is described by the following Kraus matrices:\n",
    "\n",
    "A0 = $\\begin{bmatrix}1 & 0 \\\\ 0  &  \\sqrt{1-ampl}\\end{bmatrix}$; \n",
    "A1 = $\\begin{bmatrix} 0 & \\sqrt{ampl} \\\\ 0  &  0\\end{bmatrix}$ \n",
    "\n",
    "Test case: create an amplitude damping channel with high amplitude (e.g. 25%) associated with X gates.\n",
    "\n",
    "Verify that the measure probability (no readout noise) matches expectation: 25% 0, 75% 1 (an ideal X gate produces 100% 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'{\"gate_noise\": [{\"gate_name\": \"X\", \"register_location\": [\"0\"], \"noise_channels\": [{\"matrix\": [[[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.8660254037844386, 0.0]]], [[[0.0, 0.0], [0.5, 0.0]], [[0.0, 0.0], [0.0, 0.0]]]]}]}], \"bit_order\": \"MSB\"}'"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "AD = lambda l: np.array([np.array([[1.0 , 0.0],[0.0, np.sqrt(1.0 - l)]], dtype = complex), np.array([[0.0, np.sqrt(l)], [0.0, 0.0]], dtype = complex)])\n",
    "ad_ampl = 0.25\n",
    "amplitude_damping = AD(0.25)\n",
    "# Create a simple noise model that only contains an amplitude damping channel on X gate.\n",
    "xNoise = GateNoise(\"X\", [\"0\"], [ChannelKrausOp(amplitude_damping)])\n",
    "noiseModel = NoiseModel([xNoise])\n",
    "json.dumps(noiseModel.to_dict())\n",
    "\n",
    "# Validate the simulation results: ~25% |0>, ~75% |1> measure distribution."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multi-qubit Kraus channels\n",
    "\n",
    "When specifying Kraus channels affecting multiple qubits, the 'MSB' and 'LSB' convention will help remove any ambiguity in the Kraus matrix representation.\n",
    "\n",
    "This is important for unbalanced channels (e.g. the amount of IX depolarizing is different from that of XI).\n",
    "\n",
    "Each simulator will have each own bit-order convention (for internal storage of the state-vector/density matrix); hence we need to validate the representation based on the effect on measurement results.\n",
    "\n",
    "In the following example, we create a 2-q noise channel that only acts on the first qubit (Q0).\n",
    "Depending on the LSB/MSB convention, the kron order will be different (when creating the Kraus matrices). \n",
    "We then feed the noise model (in both LSB/MSB formats) to the simulator and make sure that it will indeed only induce decoherent effect on that qubit (Q0)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a depolarizing channel on 2-q Hilbert space but only acting\n",
    "# on the first qubit (q0) according to a specific bit-ordering convention (LSB/MSB)\n",
    "def single_qubit_depol_channel_2q(gamma, order = 'MSB'):\n",
    "    sx = np.array([[0,1],[1,0]])\n",
    "    sy = np.array([[0,-1j],[1j,0]])\n",
    "    sz = np.array([[1,0],[0,-1]])\n",
    "    s0 = np.eye(2)\n",
    "    \n",
    "    if order == 'MSB':\n",
    "        # MSB: q1q0\n",
    "        return np.array((np.sqrt(1-gamma)*np.kron(s0, s0), np.sqrt(gamma/3.0)*np.kron(s0, sx), np.sqrt(gamma/3.0)*np.kron(s0, sy), np.sqrt(gamma/3.0)*np.kron(s0, sz)))\n",
    "    else:\n",
    "        # LSB: q0q1\n",
    "        # Swap the kron order:\n",
    "        return np.array((np.sqrt(1-gamma)*np.kron(s0, s0), np.sqrt(gamma/3.0)*np.kron(sx, s0), np.sqrt(gamma/3.0)*np.kron(sy, s0), np.sqrt(gamma/3.0)*np.kron(sz, s0)))      "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"gate_noise\": [{\"gate_name\": \"CNOT\", \"register_location\": [\"0\", \"1\"], \"noise_channels\": [{\"matrix\": [[[[0.99498743710662, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.99498743710662, 0.0], [0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0], [0.99498743710662, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.99498743710662, 0.0]]], [[[0.0, 0.0], [0.05773502691896258, 0.0], [0.0, 0.0], [0.0, 0.0]], [[0.05773502691896258, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.05773502691896258, 0.0]], [[0.0, 0.0], [0.0, 0.0], [0.05773502691896258, 0.0], [0.0, 0.0]]], [[[0.0, 0.0], [0.0, -0.05773502691896258], [0.0, 0.0], [0.0, 0.0]], [[0.0, 0.05773502691896258], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, -0.05773502691896258]], [[0.0, 0.0], [0.0, 0.0], [0.0, 0.05773502691896258], [0.0, 0.0]]], [[[0.05773502691896258, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [-0.05773502691896258, 0.0], [0.0, 0.0], [-0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0], [0.05773502691896258, 0.0], [0.0, 0.0]], [[0.0, 0.0], [-0.0, 0.0], [0.0, 0.0], [-0.05773502691896258, 0.0]]]]}]}], \"bit_order\": \"MSB\"}\n"
     ]
    }
   ],
   "source": [
    "gamma = 0.01                                        \n",
    "depol_msb = single_qubit_depol_channel_2q(gamma, 'MSB')\n",
    "cxNoise_msb = GateNoise(\"CNOT\", [\"0\", \"1\"], [ChannelKrausOp(depol_msb)])\n",
    "noiseModel_msb = NoiseModel([cxNoise_msb], 'MSB')\n",
    "noise_json_msb = json.dumps(noiseModel_msb.to_dict())\n",
    "print(noise_json_msb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"gate_noise\": [{\"gate_name\": \"CNOT\", \"register_location\": [\"0\", \"1\"], \"noise_channels\": [{\"matrix\": [[[[0.99498743710662, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.99498743710662, 0.0], [0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0], [0.99498743710662, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.99498743710662, 0.0]]], [[[0.0, 0.0], [0.0, 0.0], [0.05773502691896258, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.05773502691896258, 0.0]], [[0.05773502691896258, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.05773502691896258, 0.0], [0.0, 0.0], [0.0, 0.0]]], [[[0.0, 0.0], [0.0, 0.0], [0.0, -0.05773502691896258], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, -0.05773502691896258]], [[0.0, 0.05773502691896258], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.05773502691896258], [0.0, 0.0], [0.0, 0.0]]], [[[0.05773502691896258, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.05773502691896258, 0.0], [0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0], [-0.05773502691896258, 0.0], [-0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0], [-0.0, 0.0], [-0.05773502691896258, 0.0]]]]}]}], \"bit_order\": \"LSB\"}\n"
     ]
    }
   ],
   "source": [
    "gamma = 0.01                                        \n",
    "depol_lsb = single_qubit_depol_channel_2q(gamma, 'LSB')\n",
    "cxNoise_lsb = GateNoise(\"CNOT\", [\"0\", \"1\"], [ChannelKrausOp(depol_lsb)])\n",
    "noiseModel_lsb = NoiseModel([cxNoise_lsb], 'LSB')\n",
    "noise_json_lsb = json.dumps(noiseModel_lsb.to_dict())\n",
    "print(noise_json_lsb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Validation\n",
    "\n",
    "The validation here is to run a circuit that only contains a CNOT gate on different noisy simulators.\n",
    "The noise model is supplied in either LSB or MSB convention.\n",
    "Let the simulator perform the noisy simulation then perform measurements.\n",
    "\n",
    "There are two checks:\n",
    "\n",
    "(1) The resulting density matrices must be equal when using the MSB or LSB JSON files. \n",
    "\n",
    "That is to say, the convention in which the noise model is specified doesn't change the semantics of the noise model (producing the same results)\n",
    "\n",
    "(2) All the simulators (regardless of their internal MSB/LSB convention) must only show noisy results on qubit Q0.\n",
    "\n",
    "i.e. checking that we map the noise model's LSB/MSB to the simulator's LSB/MSB correctly. \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validate Readout Errors\n",
    "\n",
    "\n",
    "Specifying readout probabilities (no gate noises) and validate the readout distribution.\n",
    "\n",
    "Example: P(1|0) = 0.1; P(0|1) = 0.2\n",
    "\n",
    "- Run Identity circuit => get 90% 1, 10% 1\n",
    "\n",
    "- Run X gate circuit => get 20% 0, 80% 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"gate_noise\": [], \"bit_order\": \"MSB\", \"readout_errors\": [{\"register_location\": \"0\", \"prob_meas0_prep1\": 0.2, \"prob_meas1_prep0\": 0.1}]}\n"
     ]
    }
   ],
   "source": [
    "prep_0_meas_1 = 0.1  \n",
    "prep_1_meas_0 = 0.2                                      \n",
    "# Readout error on qubit 0\n",
    "ro_error = ReadoutError(\"0\", prep_1_meas_0, prep_0_meas_1)\n",
    "# No gate noises, just readout errors for validation.\n",
    "noiseModel = NoiseModel([], 'MSB', [ro_error])\n",
    "noise_json = json.dumps(noiseModel.to_dict())\n",
    "print(noise_json)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.4 64-bit ('base': conda)",
   "language": "python",
   "name": "python37464bitbasecondaa69d47e4607a405aa95585b12800a086"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
